import torch
import torchvision
import torch.nn as nn
import torch.nn.functional as F
from pytorch_pretrained import BertModel,BertTokenizer

class Config(object):
    """配置参数"""
    def __init__(self, dataset):
       #模型名称
       self.model_name = "BruceBertCNN"
       #训练集
       self.train_path = dataset + '/data/train.txt'
       #校验集
       self.dev_path = dataset + '/data/dev.txt'
       #测试集
       self.test_path = dataset + '/data/test.txt'
       # dataset  train/dev/test保存为pkl的路径
       self.datasetpkl = dataset + '/data/dataset.pkl'
       #类别列表
       self.class_list = [x.strip() for x in open(dataset + '/data/class.txt').readlines()]

       #模型保存路径
       self.saved_path = dataset + '/saved_dict/' + self.model_name + '.ckpt'
       #运行设备
       self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
       print('device:', self.device)
       self.require_improvement = 500
       #类别数
       self.num_classes = len(self.class_list)
       # epoch数
       self.num_epochs = 3
       # 批量大小
       self.batch_size = 128
       # 序列长度
       self.pad_size = 32
       # 学习率
       self.learning_rate = 1e-5
       # 预训练位置
       self.bert_path = './bert_pretrain'
       # bert的tokenizer
       self.tokenizer = BertTokenizer.from_pretrained(self.bert_path)
       # bert的隐藏层数量
       self.hidden_size = 768

       #TextCNN
       # 卷积核尺寸
       self.filter_size = [2,3,4]
       # 卷积核数量
       self.num_filters = 256
       # dropout率
       self.dropout = 0.1


class Model(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.bert = BertModel.from_pretrained(config.bert_path)
        for param in self.bert.parameters():
            param.requires_grad = True

        self.convs = nn.ModuleList(
            [nn.Conv2d(in_channels=1, out_channels=config.num_filters, kernel_size=(size, config.hidden_size)) for size in config.filter_size]
        )
        self.dropout = nn.Dropout(config.dropout)
        self.fc = nn.Linear(config.num_filters * len(config.filter_size), config.num_classes)

    def conv_and_pool(self, x, conv):
        x = conv(x) # shape: 128, 256, 31, 1
        x = F.relu(x)
        x = x.squeeze(3) # shape: 128,256,31
        size = x.size(2)
        x = F.max_pool1d(x, size)
        x = x.squeeze(2)
        return x


    def forward(self, x):
        # x : [ids, seq_len, mask]    input_ids: word_to_ids / id_to_words
        context = x[0]
        mask = x[2]
        encoded_layers, pooled_output = self.bert(context, attention_mask=mask, output_all_encoded_layers=False)
        out = encoded_layers.unsqueeze(1)
        out = torch.cat([self.conv_and_pool(out, conv) for conv in self.convs], 1)
        out = self.dropout(out)
        out = self.fc(out)
        return out
